const WebSocket = require('ws')
const webSocketStream = require('websocket-stream/stream')
const ffmpeg = require('fluent-ffmpeg')

// 建立WebSocket服务
const wss = new WebSocket.Server({ port: 9527, perMessageDeflate: false })

console.log('WebSocket FLV服务器已启动，监听端口 9527')
console.log('连接格式: ws://localhost:9527/rtsp://your-rtsp-url')

// 监听连接
wss.on('connection', handleConnection)

// 连接时触发事件
function handleConnection (ws, req) {
  console.log('新的WebSocket连接:', req.url)
  
  // 获取前端请求的流地址（前端websocket连接时后面带上流地址）
  const url = req.url.slice(1)
  
  // 如果是空URL，说明是服务器可用性检查，直接关闭连接
  if (!url) {
    console.log('服务器可用性检查连接')
    ws.close(1000, '服务器可用')
    return
  }
  
  if (!url.startsWith('rtsp://')) {
    console.log('缺少有效的RTSP URL参数:', url)
    ws.close(1000, '缺少有效的RTSP URL参数')
    return
  }
  
  console.log('处理RTSP流:', url)
  
  // 传入连接的ws客户端 实例化一个流
  const stream = webSocketStream(ws, { binary: true })
  
  // 通过ffmpeg命令 对实时流进行格式转换 输出flv格式
  const ffmpegCommand = ffmpeg(url)
    .addInputOption('-rtsp_transport', 'tcp')  // 使用TCP传输
    .addInputOption('-rtsp_flags', 'prefer_tcp')  // 优先使用TCP
    .addInputOption('-analyzeduration', '2000000')  // 增加分析时间
    .addInputOption('-probesize', '2000000')        // 增加探测大小
    .addInputOption('-max_delay', '2000000')        // 增加最大延迟容忍
    .addInputOption('-fflags', '+genpts+igndts')    // 生成时间戳，忽略DTS
    .addInputOption('-avoid_negative_ts', 'make_zero')
    .addInputOption('-timeout', '5000000')          // 设置超时时间5秒
    .on('start', function () { 
      console.log('Stream started for:', url) 
    })
    .on('codecData', function () { 
      console.log('Stream codecData for:', url) 
    })
    .on('error', function (err) {
      console.log('An error occured for', url, ':', err.message)
      // 不要立即结束流，给一些重试机会
      setTimeout(() => {
        if (ws.readyState === 1) { // WebSocket仍然连接
          console.log('Attempting to restart stream for:', url)
          try {
            // 使用更温和的方式终止FFmpeg进程
            ffmpegCommand.kill('SIGINT')
            setTimeout(() => {
              if (!ffmpegCommand.killed) {
                ffmpegCommand.kill('SIGTERM')
              }
              // 重新启动FFmpeg
              setTimeout(() => {
                if (ws.readyState === 1) {
                  console.log('Restarting FFmpeg for:', url)
                  ffmpegCommand.pipe(stream)
                }
              }, 1000)
            }, 2000)
          } catch (e) {
            console.log('Failed to restart stream:', e.message)
            stream.end()
          }
        } else {
          stream.end()
        }
      }, 2000)
    })
    .on('end', function () {
      console.log('Stream end for:', url)
      // 不要立即结束，给重连机会
      setTimeout(() => {
        if (ws.readyState === 1) {
          console.log('Stream ended, attempting restart for:', url)
          try {
            ffmpegCommand.pipe(stream)
          } catch (e) {
            console.log('Failed to restart after end:', e.message)
            stream.end()
          }
        } else {
          stream.end()
        }
      }, 1000)
    })
    .outputFormat('flv')
    .videoCodec('copy')
    .audioCodec('copy')
    // .outputFormat('flv').videoCodec('copy').noAudio() // 取消音频输出

  stream.on('close', function () {
    console.log('Stream closed for:', url)
    try {
      // 使用更温和的方式终止FFmpeg进程
      ffmpegCommand.kill('SIGINT')
      setTimeout(() => {
        if (!ffmpegCommand.killed) {
          ffmpegCommand.kill('SIGTERM')
        }
      }, 2000)
    } catch (e) {
      console.log('Failed to kill FFmpeg process:', e.message)
    }
  })

  stream.on('error', function (error) {
    console.log('Stream error for:', url, error.message)
    try {
      // 使用更温和的方式终止FFmpeg进程
      ffmpegCommand.kill('SIGINT')
      setTimeout(() => {
        if (!ffmpegCommand.killed) {
          ffmpegCommand.kill('SIGTERM')
        }
      }, 2000)
    } catch (e) {
      console.log('Failed to kill FFmpeg process on stream error:', e.message)
    }
  })

  try {
    // 执行命令 传输到实例流中返回给客户端
    ffmpegCommand.pipe(stream)
    console.log('FFmpeg command piped to stream for:', url)
  } catch (error) {
    console.log('FFmpeg启动失败:', error)
    stream.end()
  }
} 